# -*- coding: utf-8 -*-
"""
项目名称：NoSql_kafka_project
文件名称：check_environment.py
创建时间：2025-06-05 17:04:07

系统用户：Administrator
作　　者：無以菱
联系邮箱：huangjing2001.guet@qq.com
功能描述：
- 环境检查脚本
- 检查系统环境是否满足项目运行要求，包括Java版本、PySpark版本、Kafka、HBase和Redis连接等
"""
import os
import sys
import subprocess
import importlib
import socket
import json
import time
from importlib.metadata import version as pkg_version

# 颜色代码
GREEN = "\033[32m"
YELLOW = "\033[33m"
RED = "\033[31m"
BLUE = "\033[34m"
RESET = "\033[0m"
BOLD = "\033[1m"

# 自定义JDK路径
CUSTOM_JDK_PATH = "/usr/local/jdk"
BASHRC_PATH = "/home/wyl/.bashrc"

def print_header(message):
    """打印带颜色的标题"""
    print(f"\n{BOLD}{BLUE}=== {message} ==={RESET}")

def print_success(message):
    """打印成功消息"""
    print(f"{GREEN}✓ {message}{RESET}")

def print_warning(message):
    """打印警告消息"""
    print(f"{YELLOW}⚠ {message}{RESET}")

def print_error(message):
    """打印错误消息"""
    print(f"{RED}✗ {message}{RESET}")

def print_info(message):
    """打印信息消息"""
    print(f"{BLUE}ℹ {message}{RESET}")

def check_java_version():
    """检查Java版本"""
    print_header("检查Java版本")

    try:
        # 首先检查自定义JDK路径
        if os.path.exists(CUSTOM_JDK_PATH):
            print_info(f"检测到自定义JDK路径: {CUSTOM_JDK_PATH}")
            java_bin = os.path.join(CUSTOM_JDK_PATH, "bin", "java")
            if os.path.exists(java_bin):
                try:
                    # 使用自定义路径执行java -version
                    result = subprocess.run([java_bin, '-version'],
                                            stdout=subprocess.PIPE,
                                            stderr=subprocess.PIPE,
                                            text=True)
                    output = result.stderr if result.stderr else result.stdout
                    print_success(f"使用自定义JDK路径: {java_bin}")
                except Exception as e:
                    print_error(f"使用自定义JDK路径出错: {str(e)}")
                    output = None
            else:
                print_warning(f"自定义JDK路径不包含java可执行文件: {java_bin}")
                output = None
        else:
            print_warning(f"自定义JDK路径不存在: {CUSTOM_JDK_PATH}，将使用系统默认Java")
            # 使用系统默认Java
            try:
                result = subprocess.run(['java', '-version'],
                                        stdout=subprocess.PIPE,
                                        stderr=subprocess.PIPE,
                                        text=True)
                output = result.stderr if result.stderr else result.stdout
            except Exception as e:
                print_error(f"使用系统默认Java出错: {str(e)}")
                output = None

        if output and ("openjdk version" in output or "java version" in output):
            # 提取版本号
            version_line = output.split('\n')[0]

            if '"1.8.' in version_line:
                java_version = 8
                print_success(f"检测到Java版本: Java {java_version} ({version_line.strip()})")
            elif '"11.' in version_line:
                java_version = 11
                print_success(f"检测到Java版本: Java {java_version} ({version_line.strip()})")
            elif '"17.' in version_line:
                java_version = 17
                print_success(f"检测到Java版本: Java {java_version} ({version_line.strip()})")
            else:
                # 尝试从字符串中提取版本号
                import re
                match = re.search(r'version "([^"]+)"', version_line)
                if match:
                    version_str = match.group(1)
                    if version_str.startswith("1."):
                        java_version = int(version_str.split(".")[1])
                    else:
                        java_version = int(version_str.split(".")[0])

                    print_success(f"检测到Java版本: Java {java_version} ({version_line.strip()})")
                else:
                    java_version = None
                    print_warning(f"无法确定Java版本: {version_line.strip()}")

            return java_version
        else:
            print_error("无法获取Java版本信息")
            return None
    except FileNotFoundError:
        print_error("未安装Java或Java不在PATH中")
        return None
    except Exception as e:
        print_error(f"检查Java版本时出错: {str(e)}")
        return None

def check_java_home():
    """检查JAVA_HOME环境变量"""
    print_header("检查JAVA_HOME环境变量")

    java_home = os.environ.get('JAVA_HOME')

    if java_home:
        print_success(f"JAVA_HOME已设置: {java_home}")
        if os.path.exists(java_home):
            print_success(f"JAVA_HOME路径存在")
            java_bin = os.path.join(java_home, "bin", "java")
            if os.path.exists(java_bin):
                print_success(f"JAVA_HOME/bin/java存在: {java_bin}")
            else:
                print_warning(f"JAVA_HOME/bin/java不存在: {java_bin}")
        else:
            print_warning(f"JAVA_HOME路径不存在: {java_home}")
    else:
        print_warning("JAVA_HOME环境变量未设置")

        # 检查.bashrc文件中是否有JAVA_HOME设置
        if os.path.exists(BASHRC_PATH):
            try:
                with open(BASHRC_PATH, 'r') as f:
                    bashrc_content = f.read()
                    if "JAVA_HOME" in bashrc_content:
                        print_info(f"在{BASHRC_PATH}中找到JAVA_HOME设置，但可能未生效")
                        print_info("请运行 'source ~/.bashrc' 使环境变量生效")
                    else:
                        print_info(f"在{BASHRC_PATH}中未找到JAVA_HOME设置")
            except Exception as e:
                print_error(f"读取{BASHRC_PATH}时出错: {str(e)}")

    return java_home

def check_pyspark_version():
    """检查PySpark版本"""
    print_header("检查PySpark版本")

    try:
        # 检查PySpark是否已安装
        try:
            pyspark_version = pkg_version('pyspark')
            print_success(f"检测到PySpark版本: {pyspark_version}")
            return pyspark_version
        except Exception:
            try:
                # 尝试使用pip命令获取版本
                result = subprocess.run(['pip', 'show', 'pyspark'],
                                        stdout=subprocess.PIPE,
                                        stderr=subprocess.PIPE,
                                        text=True)
                if result.returncode == 0:
                    for line in result.stdout.split('\n'):
                        if line.startswith('Version:'):
                            pyspark_version = line.split(':')[1].strip()
                            print_success(f"检测到PySpark版本: {pyspark_version}")
                            return pyspark_version
                print_error("PySpark已安装但无法确定版本")
                return None
            except Exception:
                print_error("未安装PySpark")
                return None
    except Exception as e:
        print_error(f"检查PySpark版本时出错: {str(e)}")
        return None

def check_version_compatibility(java_version, pyspark_version):
    """检查Java和PySpark版本兼容性"""
    print_header("检查版本兼容性")

    if java_version is None or pyspark_version is None:
        print_error("无法检查版本兼容性，缺少Java或PySpark版本信息")
        return False

    # 版本兼容性表
    compatibility = {
        '4.0': {'min_java': 17, 'max_java': None},
        '3.5': {'min_java': 8, 'max_java': 17},
        '3.4': {'min_java': 8, 'max_java': 17},
        '3.3': {'min_java': 8, 'max_java': 17},
        '3.2': {'min_java': 8, 'max_java': 11},
        '3.1': {'min_java': 8, 'max_java': 11},
        '3.0': {'min_java': 8, 'max_java': 11},
        '2.4': {'min_java': 8, 'max_java': 8}
    }

    # 提取主要版本号
    major_minor = '.'.join(pyspark_version.split('.')[:2])

    if major_minor not in compatibility:
        print_warning(f"未知的PySpark版本 {pyspark_version}，无法确定兼容性")
        return False

    min_java = compatibility[major_minor]['min_java']
    max_java = compatibility[major_minor]['max_java']

    if max_java is None:
        is_compatible = java_version >= min_java
        if is_compatible:
            print_success(f"版本兼容: Java {java_version} 兼容 PySpark {pyspark_version} (需要 Java {min_java}+)")
        else:
            print_error(f"版本不兼容: PySpark {pyspark_version} 需要 Java {min_java}+，但当前是 Java {java_version}")
    else:
        is_compatible = min_java <= java_version <= max_java
        if is_compatible:
            print_success(f"版本兼容: Java {java_version} 兼容 PySpark {pyspark_version} (需要 Java {min_java}-{max_java})")
        else:
            print_error(f"版本不兼容: PySpark {pyspark_version} 需要 Java {min_java}-{max_java}，但当前是 Java {java_version}")

    return is_compatible

def check_kafka_connection():
    """检查Kafka连接"""
    print_header("检查Kafka连接")

    try:
        # 尝试导入kafka-python库
        try:
            from kafka import KafkaProducer
            print_success("kafka-python库已安装")
        except ImportError:
            print_error("未安装kafka-python库")
            return False

        # 尝试从config.py导入Kafka配置
        try:
            sys.path.append(os.getcwd())
            from config import KAFKA_CONFIG
            bootstrap_servers = KAFKA_CONFIG['bootstrap_servers']
            print_success(f"已加载Kafka配置: {bootstrap_servers}")
        except (ImportError, KeyError):
            bootstrap_servers = ['localhost:9092']
            print_warning(f"无法从config.py加载Kafka配置，使用默认值: {bootstrap_servers}")

        # 尝试连接Kafka
        print_info(f"尝试连接Kafka服务器: {bootstrap_servers}")

        # 对每个服务器尝试建立TCP连接
        connected = False
        for server in bootstrap_servers:
            host, port = server.split(':')
            port = int(port)

            try:
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.settimeout(5)
                result = sock.connect_ex((host, port))
                sock.close()

                if result == 0:
                    print_success(f"成功连接到Kafka服务器: {server}")
                    connected = True
                    break
                else:
                    print_error(f"无法连接到Kafka服务器: {server}")
            except Exception as e:
                print_error(f"连接Kafka服务器 {server} 时出错: {str(e)}")

        if not connected:
            print_error("无法连接到任何Kafka服务器")
            return False

        # 尝试创建KafkaProducer
        try:
            producer = KafkaProducer(
                bootstrap_servers=bootstrap_servers,
                client_id='environment-check',
                request_timeout_ms=5000,
                api_version=(0, 10, 0)
            )
            producer.close()
            print_success("成功创建Kafka生产者")
            return True
        except Exception as e:
            print_error(f"创建Kafka生产者时出错: {str(e)}")
            return False
    except Exception as e:
        print_error(f"检查Kafka连接时出错: {str(e)}")
        return False

def check_hbase_connection():
    """检查HBase连接"""
    print_header("检查HBase连接")

    try:
        # 尝试导入happybase库
        try:
            import happybase
            print_success("happybase库已安装")
        except ImportError:
            print_error("未安装happybase库")
            return False

        # 尝试从config.py导入HBase配置
        try:
            sys.path.append(os.getcwd())
            from config import HBASE_CONFIG
            host = HBASE_CONFIG.get('host', 'localhost')
            port = HBASE_CONFIG.get('port', 9090)
            print_success(f"已加载HBase配置: {host}:{port}")
        except (ImportError, AttributeError):
            host = 'localhost'
            port = 9090
            print_warning(f"无法从config.py加载HBase配置，使用默认值: {host}:{port}")

        # 尝试连接HBase
        print_info(f"尝试连接HBase服务器: {host}:{port}")

        # 首先尝试建立TCP连接
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(5)
            result = sock.connect_ex((host, port))
            sock.close()

            if result == 0:
                print_success(f"成功连接到HBase Thrift服务器: {host}:{port}")
            else:
                print_error(f"无法连接到HBase Thrift服务器: {host}:{port}")
                return False
        except Exception as e:
            print_error(f"连接HBase服务器时出错: {str(e)}")
            return False

        # 尝试创建HBase连接
        try:
            connection = happybase.Connection(host=host, port=port, timeout=5000)
            tables = connection.tables()
            connection.close()
            print_success(f"成功连接到HBase，可用表: {tables}")
            return True
        except Exception as e:
            print_error(f"创建HBase连接时出错: {str(e)}")
            return False
    except Exception as e:
        print_error(f"检查HBase连接时出错: {str(e)}")
        return False

def check_redis_connection():
    """检查Redis连接"""
    print_header("检查Redis连接")

    try:
        # 尝试导入redis库
        try:
            import redis
            print_success("redis库已安装")
        except ImportError:
            print_error("未安装redis库")
            return False

        # 尝试从config.py导入Redis配置
        try:
            sys.path.append(os.getcwd())
            from config import REDIS_CONFIG
            host = REDIS_CONFIG.get('host', 'localhost')
            port = REDIS_CONFIG.get('port', 6379)
            db = REDIS_CONFIG.get('db', 0)
            print_success(f"已加载Redis配置: {host}:{port}/{db}")
        except (ImportError, AttributeError):
            host = 'localhost'
            port = 6379
            db = 0
            print_warning(f"无法从config.py加载Redis配置，使用默认值: {host}:{port}/{db}")

        # 尝试连接Redis
        print_info(f"尝试连接Redis服务器: {host}:{port}/{db}")

        # 首先尝试建立TCP连接
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(5)
            result = sock.connect_ex((host, port))
            sock.close()

            if result == 0:
                print_success(f"成功连接到Redis服务器: {host}:{port}")
            else:
                print_error(f"无法连接到Redis服务器: {host}:{port}")
                return False
        except Exception as e:
            print_error(f"连接Redis服务器时出错: {str(e)}")
            return False

        # 尝试创建Redis连接
        try:
            r = redis.Redis(host=host, port=port, db=db)
            r.ping()
            print_success("成功连接到Redis")
            return True
        except Exception as e:
            print_error(f"创建Redis连接时出错: {str(e)}")
            return False
    except Exception as e:
        print_error(f"检查Redis连接时出错: {str(e)}")
        return False

def check_python_packages():
    """检查Python包依赖"""
    print_header("检查Python包依赖")

    required_packages = {
        'pyspark': '3.3.3',
        'happybase': None,
        'kafka-python': None,
        'redis': None,
        'flask': None,
        'pandas': None
    }

    all_installed = True

    for package, required_version in required_packages.items():
        try:
            installed_version = pkg_version(package)
            if required_version and installed_version != required_version:
                print_warning(f"{package}: 已安装 {installed_version}，但需要 {required_version}")
                all_installed = False
            else:
                print_success(f"{package}: {installed_version}")
        except Exception:
            try:
                # 尝试导入包
                importlib.import_module(package)
                print_warning(f"{package}: 已安装，但无法确定版本")
            except ImportError:
                print_error(f"{package}: 未安装")
                all_installed = False

    return all_installed


def main():
    """主函数"""
    print(f"{BOLD}{BLUE}===================================={RESET}")
    print(f"{BOLD}{BLUE}  实时支付数据分析系统 - 环境检查  {RESET}")
    print(f"{BOLD}{BLUE}===================================={RESET}")

    # 检查Java版本
    java_version = check_java_version()

    # 检查JAVA_HOME环境变量
    java_home = check_java_home()

    # 检查PySpark版本
    pyspark_version = check_pyspark_version()

    # 检查版本兼容性
    is_compatible = check_version_compatibility(java_version, pyspark_version)

    # 检查Python包依赖
    packages_ok = check_python_packages()

    # 检查Kafka连接
    kafka_ok = check_kafka_connection()

    # 检查HBase连接
    hbase_ok = check_hbase_connection()

    # 检查Redis连接
    redis_ok = check_redis_connection()


    # 总结
    print_header("检查结果总结")

    all_ok = is_compatible and packages_ok and kafka_ok and hbase_ok and redis_ok

    if all_ok:
        print_success("所有检查都通过了，环境配置正确")
    else:
        print_warning("环境检查发现以下问题:")

        if not is_compatible:
            print_error("- Java和PySpark版本不兼容")

        if not packages_ok:
            print_error("- 缺少必要的Python包")

        if not kafka_ok:
            print_error("- Kafka连接失败")

        if not hbase_ok:
            print_error("- HBase连接失败")

        if not redis_ok:
            print_error("- Redis连接失败")

        print_info("\n解决方案:")

        if not is_compatible:
            if java_version and pyspark_version:
                major_pyspark = '.'.join(pyspark_version.split('.')[:2])
                if major_pyspark == '4.0' and java_version < 17:
                    print_info("- 将PySpark降级到与Java兼容的版本: pip install pyspark==3.3.3")
                elif java_version == 8 and major_pyspark not in ['2.4', '3.0', '3.1', '3.2', '3.3', '3.4', '3.5']:
                    print_info("- 将PySpark降级到与Java 8兼容的版本: pip install pyspark==3.3.3")
                else:
                    print_info(f"- 升级Java到版本17: sudo apt install openjdk-17-jdk")

            # 添加JAVA_HOME设置建议
            if not java_home:
                print_info(f"- 设置JAVA_HOME环境变量: export JAVA_HOME={CUSTOM_JDK_PATH}")
                print_info(f"- 将以下行添加到{BASHRC_PATH}:")
                print_info(f"  export JAVA_HOME={CUSTOM_JDK_PATH}")
                print_info(f"  export PATH=$JAVA_HOME/bin:$PATH")

        if not packages_ok:
            print_info("- 安装缺少的Python包: pip install -r requirements.txt")

        if not kafka_ok:
            print_info("- 检查Kafka服务是否运行")
            print_info("- 检查config.py中的Kafka配置是否正确")
            print_info("- 确保网络连接正常，防火墙允许连接")

        if not hbase_ok:
            print_info("- 检查HBase Thrift服务是否运行")
            print_info("- 检查config.py中的HBase配置是否正确")
            print_info("- 确保网络连接正常，防火墙允许连接")

        if not redis_ok:
            print_info("- 检查Redis服务是否运行")
            print_info("- 检查config.py中的Redis配置是否正确")
            print_info("- 确保网络连接正常，防火墙允许连接")

    return 0 if all_ok else 1

if __name__ == "__main__":
    sys.exit(main())
